"""Simulation tool functions"""

import numpy
from kuai.ff import ForceFieldItem

from kuai.kuaiext import SIZE_OF_REAL_NUMBER, SIZE_OF_ATOM_INDEX
if SIZE_OF_REAL_NUMBER == 8:
    REAL_TYPE = numpy.float64
elif SIZE_OF_REAL_NUMBER == 4:
    REAL_TYPE = numpy.float32
else:
    assert False

if SIZE_OF_ATOM_INDEX == 8:
    ATOM_INDEX = numpy.int64
elif SIZE_OF_ATOM_INDEX == 4:
    ATOM_INDEX = numpy.int32
elif SIZE_OF_ATOM_INDEX == 2:
    ATOM_INDEX = numpy.int16
else:
    assert False

def collect_coords(mol):
    result = []
    for i in mol.atoms:
        result.append(i.coords.x)
        result.append(i.coords.y)
        result.append(i.coords.z)
    return numpy.array(result, dtype=REAL_TYPE)

def collect_weights(mol, ff=None):
    result = []
    for i in mol.atoms:
        item = ForceFieldItem(FF_ATYPE, i.type)
        if ff is not None and item in ff: # Try to find weight in force field
            weight = ff[item][1]
        else:   # Use default if it is not existed
            weight = i.weight()
        result += [weight] * 3
    return numpy.array(result, dtype=REAL_TYPE)

def collect_k(ff):
    nk = 0
    result = []
    for k, v in ff.iteritems():
        v.index = nk
        nk += len(v)
        result += v
    return numpy.array(result, dtype=REAL_TYPE)

def setup_energy_terms(mol, ff, term):
    func = ffman[term]
    result = []
    atoms = func.list_atoms(mol)
    rank = func.rank()
    for i in range(0, len(atoms), rank):
        atom_index = [mol.index(i) * 3 for i in atoms[i:i+rank]]
        atom_types = [i.type for i in atoms[i:i+rank]]
        func.uniform(atom_types, None, atom_index)
        item = ForceFieldItem(term, atom_types)
        if item in ff:
            result += atom_index
            result.append(ff[item].index)
    return numpy.array(result, dtype=ATOM_INDEX) 

